package com.sz.admin.ai.factory.embeddingFactory.Handler.impl;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingOptions;
import com.sz.admin.ai.factory.embeddingFactory.Handler.EmbeddingHandler;
import com.sz.admin.ai.factory.embeddingFactory.entity.EmbeddingHandlerRequestDTO;
import com.sz.admin.ai.factory.embeddingFactory.entity.EmbeddingHandlerUser;
import com.sz.admin.ai.factory.embeddingFactory.entity.EmbeddingModelEntity;
import com.sz.admin.ai.factory.embeddingFactory.entity.EmbeddingPlatformEnum;
import com.sz.admin.ai.util.PgVectorStoreUtil;
import com.sz.core.util.StringUtils;
import jakarta.annotation.Resource;
import lombok.extern.log4j.Log4j2;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @描述:
 * @Author: TangYuan
 * @Date: 2025/2/25 14:59
 */
@Component
@Log4j2
public class AliyunEmbeddingHandlerImpl implements EmbeddingHandler {

    @Resource(name = "pgAliyunJdbcTemplate")
    private JdbcTemplate jdbcTemplate;
    @Resource
    private PgVectorStoreUtil pgVectorStoreUtil;

    /**
     * 获取EmbeddingHandler
     * @return {@link EmbeddingPlatformEnum}
     */
    @Override
    public EmbeddingPlatformEnum getEmbeddingPlatformEnum() {
        return EmbeddingPlatformEnum.ALIYUN;
    }

    /**
     * 向量化数据 -> 存向量化数据
     * @param embeddingHandlerRequestDTO 参数
     * @return {@link Boolean}
     */
    @Override
    public Boolean vectorizedData(EmbeddingHandlerRequestDTO embeddingHandlerRequestDTO) {
        try {
            // 第三方应用名
            String thirdPartyName = embeddingHandlerRequestDTO.getThirdPartyName();
            // 模型名
            String modelName = StringUtils.isEmpty(embeddingHandlerRequestDTO.getModelName()) ? "text-embedding-v3" : embeddingHandlerRequestDTO.getModelName();
            // 用户信息
            EmbeddingHandlerUser embeddingHandlerUser = embeddingHandlerRequestDTO.getEmbeddingHandlerUser();
            // 向量化文本
            List<Document> documentList = this.getDocumentList(embeddingHandlerRequestDTO.getDocumentList(), thirdPartyName, embeddingHandlerUser, modelName, embeddingHandlerRequestDTO.getFileName());

            EmbeddingModelEntity modelEntity = this.getModelEntity(modelName);

            DashScopeApi dashScopeApi = new DashScopeApi(embeddingHandlerRequestDTO.getModelKey());

            DashScopeEmbeddingModel embeddingModel = new DashScopeEmbeddingModel(
                    dashScopeApi,
                    MetadataMode.ALL,
                    DashScopeEmbeddingOptions.builder()
                            .withModel(modelEntity.getModelName())
                            .withDimensions(modelEntity.getModelDimension())
                            .withTextType("document")
                            .build()
            );

            /*
             * 封装
             * 根据模型名获取VectorStore
             * 模型的维度 -> 对应数据库的名称，和数据库的维度
             */
            VectorStore vectorStore = pgVectorStoreUtil.getPgVectorStore(jdbcTemplate, embeddingModel, this.getModelEntity(modelName));

            //* 进行PG数据库向量化和添加操作
            vectorStore.add(documentList);

            return true;
        }catch (Exception e) {
            log.error(e.getMessage());
            return false;
        }
    }

    /**
     * 使用向量化数据
     * @param embeddingHandlerRequestDTO 基本参数
     * @param advisorSpec
     * @param textContent 文本内容
     */
    @Override
    public void useVectorStore(EmbeddingHandlerRequestDTO embeddingHandlerRequestDTO, Integer topK, ChatClient.AdvisorSpec advisorSpec, String textContent) {

        // 第三方应用名
        String thirdPartyName = embeddingHandlerRequestDTO.getThirdPartyName();
        // 模型名
        String modelName = StringUtils.isEmpty(embeddingHandlerRequestDTO.getModelName()) ? "text-embedding-v3" : embeddingHandlerRequestDTO.getModelName();
        // 用户信息
        EmbeddingHandlerUser embeddingHandlerUser = embeddingHandlerRequestDTO.getEmbeddingHandlerUser();

        String promptWithContext = """
                 下面是上下文信息
                ---------------------
                 {question_answer_context}
                 ---------------------
                 给定的上下文和提供的历史信息，而不是事先的知识，回复用户的意见。如果答案不在上下文中，告诉用户你不能回答这个问题。
                """;

        EmbeddingModelEntity modelEntity = this.getModelEntity(modelName);

        DashScopeApi dashScopeApi = new DashScopeApi(embeddingHandlerRequestDTO.getModelKey());

        DashScopeEmbeddingModel embeddingModel = new DashScopeEmbeddingModel(
                dashScopeApi,
                MetadataMode.ALL,
                DashScopeEmbeddingOptions.builder()
                        .withModel(modelEntity.getModelName())
                        .withDimensions(modelEntity.getModelDimension())
                        .withTextType("document")
                        .build()
        );

        //创建一个向量库连接
        VectorStore vectorStore = pgVectorStoreUtil.getPgVectorStore(jdbcTemplate, embeddingModel, this.getModelEntity(modelName));

        //设置查询中的过滤条件
        FilterExpressionBuilder builder = new FilterExpressionBuilder();
        Filter.Expression filterExpression = builder
                .and(
                        builder.and(
                                builder.eq("thirdPartyName", thirdPartyName),
                                builder.eq("userId", thirdPartyName + "-" + embeddingHandlerUser.getUserId())),
                        builder.eq("knowledgeBaseId", thirdPartyName + "-" + embeddingHandlerUser.getKnowledgeBaseId())
                ).build();

        advisorSpec.advisors(new QuestionAnswerAdvisor(
                vectorStore,
                SearchRequest.builder()
                        .query(textContent)
                        .filterExpression(filterExpression)
                        .topK(topK)
                        .build(),
                promptWithContext));
    }

    @Override
    public Integer delete(EmbeddingHandlerRequestDTO embeddingHandlerRequestDTO) {
        // 第三方应用名
        String thirdPartyName = embeddingHandlerRequestDTO.getThirdPartyName();
        // 模型名
        String modelName = StringUtils.isEmpty(embeddingHandlerRequestDTO.getModelName()) ? "text-embedding-v3" : embeddingHandlerRequestDTO.getModelName();
        // 用户信息
        EmbeddingHandlerUser embeddingHandlerUser = embeddingHandlerRequestDTO.getEmbeddingHandlerUser();

        String tableName = "public.vector_store_" + getModelEntity(modelName).getModelDimension();

        Map<String, String> map = new HashMap<>();
        map.put("thirdPartyName", thirdPartyName);
        map.put("userId", embeddingHandlerUser.getUserId());
        map.put("knowledgeBaseId", embeddingHandlerUser.getKnowledgeBaseId());
        if (StringUtils.isEmpty(embeddingHandlerRequestDTO.getFileName())) {
            map.put("fileId", embeddingHandlerRequestDTO.getFileId());
        }

        String sql = "DELETE FROM " + tableName + " WHERE metadata::jsonb @@ ?::jsonpath";
        return jdbcTemplate.update(sql, this.generateJsonPath(map));
    }
}
